import random

import numpy as np
import torch

from model import SSQNet
from tools import ReadConfigs, ReadIntArray, OneHotData, GetParamList


def predict():
    configs = ReadConfigs()
    red_file_path = configs["RED_FILE_PATH"]
    sample_nums = configs["SAMPLE_NUMS"]

    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    my_list = ReadIntArray(red_file_path)
    data_list = my_list[-11:]
    data_list = OneHotData(data_list)
    data_list = torch.tensor(data_list, dtype=torch.float32)
    data_list = data_list.unsqueeze(dim=0)

    a = torch.tensor([0])
    b = torch.tensor([1])
    list1 = []
    list2 = []
    for index in range(1, 34):
        model = SSQNet(input_size=33, hidden_size=128, num_layers=3, num_classes=2, sample_nums=sample_nums)
        model = model.to(device)
        path_name = "pth/red_predict_" + str(index) + ".pth"
        model.load_state_dict(torch.load(path_name))

        model.eval()
        pred_list = model(data_list)
        pred_target = torch.max(pred_list, 1)[1]

        if torch.eq(a, pred_target):
            list1.append(index)
        elif torch.eq(b, pred_target):
            list2.append(index)

    return list1, list2


def GetPredList(list1, list2, param_lists):
    pred_list = []
    while len(pred_list) < 8:
        tmp_num = random.randrange(1, 34, 1)
        if tmp_num not in pred_list:
            if tmp_num in list1:
                if round(np.random.uniform(0, 1), 5) > (1 - param_lists[tmp_num - 1][0]):
                    pred_list.append(tmp_num)
                else:
                    continue
            elif tmp_num in list2:
                if round(np.random.uniform(0, 1), 5) > param_lists[tmp_num - 1][1]:
                    pred_list.append(tmp_num)
                else:
                    continue
            else:
                print("error")
    pred_list.sort()
    return set(pred_list)


if __name__ == "__main__":
    list1, list2 = predict()
    param_lists = GetParamList()
    true_list = set([3, 18, 23, 24, 25, 32])

    index = 0
    while 1:
        index += 1
        pred_list = GetPredList(list1, list2, param_lists)
        print(index, pred_list)
        if true_list <= pred_list:
            break
